/* -- INCLUDES ---------------------------------------------------------------*/
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h> 
#include <fcntl.h>
#include <sys/mman.h>

#include "Obf.h"

/* -- DEFINES ----------------------------------------------------------------*/

/* -- PRIVATE VARIABLES ------------------------------------------------------*/
static char * output_dir		= NULL;
static char * original_file		= NULL;
static int number_of_files		= 0;
static int count				= 0;
static double ratio				= 0;

/* -- PRIVATE FUNCTIONS ------------------------------------------------------*/
static void usage( char * prog )
{
	printf("Usage : %s [OPTIONS]\n", prog ? prog : DEFAULT_PROG_NAME );
	printf("\t-d, --dir <dir>\toutput directory for fuzzed files (default is %s)\n", DEFAULT_FUZZ_DIR);
	printf("\t-f, --file <file>\tfile to fuzz (required)\n");
	printf("\t-c, --count\texact count of fuzzed bytes in each produced file (not used by default)\n");
	printf("\t-r, --ratio\tratio of fuzzed bytes in each produced file (in [0, 10%] default is %f)\n", DEFAULT_RATIO);
	printf("\t-n, --number\tnumber of file to produce (default is infinite)\n");
	printf("\t-h, --help\tthis help\n");

	exit(EXIT_FAILURE);
}

static void handle_options( int argc, char ** argv )
{
	struct option options[] = {
		{"dir",		required_argument,	NULL,	'd'},
		{"file",	required_argument,	NULL,	'f'},
		{"count",	required_argument,	NULL,	'c'},
		{"ratio",	required_argument,	NULL,	'r'},
		{"number",	required_argument,	NULL,	'n'},
		{"help",	no_argument,		NULL,	'h'},
		{0,			0,					0,		0}
	};

	char c;
	int index;

	while( (c = getopt_long(argc, argv, "d:f:c:r:n:h", options, &index)) != -1 ) {
		switch( c ) {
			case 'd':
				output_dir = strdup(optarg);
				break;

			case 'f':
				original_file = strdup(optarg);
				break;
				
			case 'c':
				count = atoi(optarg);
				break;

			case 'r':
				ratio = atof(optarg);
				break;

			case 'n':
				number_of_files = atoi(optarg);
				break;

			case 'h':
				usage( NULL );
				break;

			default	:
				if( isalnum(c) )
					FFATAL("[-] Unknow option \"%c\"\n", c);
				else
					FATAL("[-] Unknow option\n");
		}
	}
	
	if( output_dir == NULL )
		output_dir = DEFAULT_FUZZ_DIR;
		
	if( ratio == 0 && count == 0 )
		ratio = DEFAULT_RATIO;
}

static char * basename( char * filename )
{
	unsigned int len = strlen(filename);
	while( len >= 0 && filename[len--] != '/' );
	
	return filename + len + 1;
}

double uniform_deviate( int seed )
{
	return seed * ( 1.0 / ( RAND_MAX + 1.0 ) );
}

int bounded_rand( int N )
{
	return uniform_deviate ( rand() ) * N;
}

static unsigned time_seed()
{
	time_t now = time( 0 );
	unsigned char *p = (unsigned char *)&now;
	unsigned seed = 0;
	size_t i;

	for ( i = 0; i < sizeof(now); i++ )
		seed = seed * ( UCHAR_MAX + 2U ) + p[i]; 

	srand( seed );
}


/* -- PUBLIC FUNCTIONS -------------------------------------------------------*/
int main( int argc, char ** argv )
{
	struct stat stats;
	int original, fuzzed;
	char fuzzed_file[FILENAME_SIZE + 1] = {0};
	unsigned int file_count = 0;
	unsigned int file_size = 0;
	
	// output_dir ??
	
	if( original_file == NULL )
		FATAL("[-] You must specify an original file with option -f\n");
	
	fd_original = open(original_file, O_RDONLY | O_DIRECT);
	if( original == -1 )
		FFATAL("[-] Cannot open original file \"%s\" : %s", original_file, strerror(errno));
		
	stats = fstat(fd_original, &stats);
	if( stats == -1 )
		PFATAL("[-] Cannot stat original file ");
	file_size = stats.st_size;
	
	original = (char *)mmap(0, (size_t)file_size, PROT_READ, MAP_PRIVATE, fd_original, 0);
	if( original == MAP_FAILED )
		PFATAL("[-] Cannot map original file ");
		
	if( number_of_files < 0 )
		FFATAL("[-] Invalid number of files : %d\n", number_of_files);
		
	if( ratio < 0 || ratio > MAX_RATIO )
		FFATAL("[-] Invalid ratio : %f\n", ratio);
		
	if( count < 0 || count > MAX_COUNT || count > file_size )
		FFATAL("[-] Invalid count of bytes : %d\n", count);

	if( ratio > 0 && count > 0 )
		FATAL("[-] You cannot specify ratio and count at the same time\n");
	else if( count > 0 )
		ratio = (double)count / (double)file_size;
		
	time_seed();

	while( number_of_files == 0 || file_count < number_of_files ) {
		if( snprintf(fuzzed_file, FILENAME_SIZE, FILENAME_FRMT, output_dir, file_count, basename(original_file)) < 0 ) {
			DEBUG("[~] cannot create file name for file %u\n", file_count);
		} else {
			fd_fuzzed = open(fuzzed_file, O_WRONLY | O_DIRECT);
			if( fd_fuzzed == -1 ) {
				DEBUG("[~] cannot create file %u\n", file_count);
			} else {
				fuzzed = (char *)mmap(0, (size_t)file_size, PROT_READ, MAP_PRIVATE, fd_fuzzed, 0);
				if( fuzzed == MAP_FAILED ) {
					DEBUG("[~] cannot map file %u\n", file_count);
				} else {
					memcpy((void *)fuzzed, (void *)original, file_size);
					for( i=0; i<count; i++ )
						fuzzed[ bounded_rand(file_size-1) ] = (char)bounded_rand(MAX_BYTE);
						
					munmap(fuzzed, file_size);
					close(fd_fuzzed);
					file_count++;
				}
			}
		}
	}

	return EXIT_SUCCESS;
}
